"""
DaoMath Quantum Mechanics (道数量子力学)
从 ±0 先天均衡构造量子理论

量子力学的 ±0 诠释:
    |ψ⟩ (量子态) - 正向 (实在)
    ⟨ψ| (对偶态) - 负向 (观测)
    测量坍缩 - 归于均衡

理论层级:
    ±0 (先天均衡)
     ↓
    Hilbert空间 H (状态空间)
     ↓
    量子态 |ψ⟩ (波函数)
     ↓
    算符 Â (可观测量)
     ↓
    测量理论 (Born规则)
     ↓
    量子演化 (Schrödinger方程)
     ↓
    纠缠与非局域性
"""

from __future__ import annotations
from typing import (
    List, Tuple, Callable, Optional, Dict, Any,
    Union, TypeVar, Generic
)
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
from enum import Enum
import numpy as np
import scipy.linalg as la
from scipy.integrate import odeint
import matplotlib.pyplot as plt

from core.axioms import InnateBalance, DualPair, ConservationLaw
from structures.number import Complex, Real
from structures.space import HilbertSpace, InnerProduct


# ============================================================================
# 基础常数
# ============================================================================

class QuantumConstants:
    """量子力学基本常数"""
    HBAR = 1.054571817e-34  # 约化普朗克常数 (J·s)
    C = 299792458.0          # 光速 (m/s)
    ELECTRON_MASS = 9.1093837015e-31  # 电子质量 (kg)
    ELEMENTARY_CHARGE = 1.602176634e-19  # 元电荷 (C)
    BOLTZMANN = 1.380649e-23  # 玻尔兹曼常数 (J/K)


# ============================================================================
# Level 1: 量子态 (Quantum State)
# ============================================================================

class QuantumState:
    """
    量子态 |ψ⟩
    
    基于 ±0 的理解:
    - 波函数表示：ψ(x) = ⟨x|ψ⟩
    - 概率诠释：|ψ(x)|² = 测量概率密度
    - 归一化：⟨ψ|ψ⟩ = 1
    """
    
    def __init__(self, state_vector: np.ndarray, is_normalized: bool = False):
        """
        Args:
            state_vector: 复数向量
            is_normalized: 是否已归一化
        """
        self.state_vector = np.array(state_vector, dtype=complex)
        
        if not is_normalized:
            self.normalize()
    
    def normalize(self) -> QuantumState:
        """归一化: ||ψ|| = 1"""
        norm = np.linalg.norm(self.state_vector)
        if norm > 1e-10:
            self.state_vector = self.state_vector / norm
        return self
    
    def dual(self) -> QuantumState:
        """对偶态 ⟨ψ| = |ψ⟩†"""
        return QuantumState(np.conj(self.state_vector), is_normalized=True)
    
    def inner_product(self, other: QuantumState) -> complex:
        """内积 ⟨ψ|φ⟩"""
        return np.vdot(self.state_vector, other.state_vector)
    
    def norm(self) -> float:
        """范数 ||ψ||"""
        return float(np.linalg.norm(self.state_vector))
    
    def probability_amplitude(self, basis_state: QuantumState) -> complex:
        """
        概率幅: ⟨φ|ψ⟩
        
        概率 = |⟨φ|ψ⟩|²
        """
        return self.inner_product(basis_state)
    
    def measurement_probability(self, basis_state: QuantumState) -> float:
        """
        测量概率 (Born规则):
        P(φ) = |⟨φ|ψ⟩|²
        """
        amplitude = self.probability_amplitude(basis_state)
        return float(np.abs(amplitude) ** 2)
    
    def expectation_value(self, operator: QuantumOperator) -> complex:
        """
        期望值: ⟨Â⟩ = ⟨ψ|Â|ψ⟩
        """
        return operator.expectation(self)
    
    def uncertainty(self, operator: QuantumOperator) -> float:
        """
        不确定度: ΔA = √(⟨Â²⟩ - ⟨Â⟩²)
        """
        A_mean = self.expectation_value(operator)
        A2_mean = self.expectation_value(operator @ operator)
        variance = A2_mean - A_mean ** 2
        return float(np.sqrt(np.abs(variance)))
    
    def __add__(self, other: QuantumState) -> QuantumState:
        """叠加态: |ψ⟩ + |φ⟩"""
        return QuantumState(
            self.state_vector + other.state_vector,
            is_normalized=False
        )
    
    def __mul__(self, scalar: complex) -> QuantumState:
        """标量乘法: c|ψ⟩"""
        return QuantumState(
            scalar * self.state_vector,
            is_normalized=False
        )
    
    def __rmul__(self, scalar: complex) -> QuantumState:
        return self.__mul__(scalar)
    
    def __repr__(self) -> str:
        if len(self.state_vector) <= 4:
            return f"|ψ⟩ = {self.state_vector}"
        return f"|ψ⟩ (dim={len(self.state_vector)})"
    
    @staticmethod
    def from_innate_balance(positive: float, negative: float) -> QuantumState:
        """
        从 ±0 构造量子态
        
        |ψ⟩ = α|0⟩ + β|1⟩
        α = √positive, β = √negative
        """
        total = positive + np.abs(negative)
        if total > 1e-10:
            alpha = np.sqrt(positive / total)
            beta = np.sqrt(np.abs(negative) / total)
            return QuantumState([alpha, beta])
        return QuantumState([1, 0])


# ============================================================================
# Level 2: 量子算符 (Quantum Operator)
# ============================================================================

class QuantumOperator:
    """
    量子算符 Â
    
    表示可观测量（Hermitian算符）或演化（Unitary算符）
    """
    
    def __init__(self, matrix: np.ndarray):
        """
        Args:
            matrix: 算符的矩阵表示
        """
        self.matrix = np.array(matrix, dtype=complex)
        self.dimension = self.matrix.shape[0]
        
        if self.matrix.shape[0] != self.matrix.shape[1]:
            raise ValueError("Operator matrix must be square")
    
    def is_hermitian(self, tol: float = 1e-10) -> bool:
        """判断是否为Hermitian算符: Â† = Â"""
        return np.allclose(self.matrix, self.matrix.conj().T, atol=tol)
    
    def is_unitary(self, tol: float = 1e-10) -> bool:
        """判断是否为Unitary算符: Û†Û = I"""
        product = self.matrix.conj().T @ self.matrix
        identity = np.eye(self.dimension)
        return np.allclose(product, identity, atol=tol)
    
    def eigendecomposition(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        本征分解: Â = Σᵢ λᵢ|λᵢ⟩⟨λᵢ|
        
        Returns:
            eigenvalues: 本征值数组
            eigenvectors: 本征向量矩阵（列向量）
        """
        eigenvalues, eigenvectors = np.linalg.eig(self.matrix)
        return eigenvalues, eigenvectors
    
    def expectation(self, state: QuantumState) -> complex:
        """期望值: ⟨ψ|Â|ψ⟩"""
        Apsi = self.matrix @ state.state_vector
        return np.vdot(state.state_vector, Apsi)
    
    def commutator(self, other: QuantumOperator) -> QuantumOperator:
        """
        对易子: [Â, B̂] = ÂB̂ - B̂Â
        """
        comm = self.matrix @ other.matrix - other.matrix @ self.matrix
        return QuantumOperator(comm)
    
    def anticommutator(self, other: QuantumOperator) -> QuantumOperator:
        """
        反对易子: {Â, B̂} = ÂB̂ + B̂Â
        """
        anticomm = self.matrix @ other.matrix + other.matrix @ self.matrix
        return QuantumOperator(anticomm)
    
    def trace(self) -> complex:
        """迹: Tr(Â) = Σᵢ ⟨i|Â|i⟩"""
        return np.trace(self.matrix)
    
    def __matmul__(self, other: Union[QuantumOperator, QuantumState]) -> Union[QuantumOperator, QuantumState]:
        """算符作用或算符乘积"""
        if isinstance(other, QuantumState):
            # Â|ψ⟩
            result_vector = self.matrix @ other.state_vector
            return QuantumState(result_vector, is_normalized=False)
        elif isinstance(other, QuantumOperator):
            # ÂB̂
            return QuantumOperator(self.matrix @ other.matrix)
        else:
            raise TypeError("Unsupported operand type")
    
    def __add__(self, other: QuantumOperator) -> QuantumOperator:
        """算符加法"""
        return QuantumOperator(self.matrix + other.matrix)
    
    def __mul__(self, scalar: complex) -> QuantumOperator:
        """标量乘法"""
        return QuantumOperator(scalar * self.matrix)
    
    def __rmul__(self, scalar: complex) -> QuantumOperator:
        return self.__mul__(scalar)
    
    def __repr__(self) -> str:
        if self.dimension <= 4:
            return f"Operator:\n{self.matrix}"
        return f"Operator (dim={self.dimension}×{self.dimension})"
    
    @staticmethod
    def identity(dimension: int) -> QuantumOperator:
        """单位算符"""
        return QuantumOperator(np.eye(dimension))
    
    @staticmethod
    def zero(dimension: int) -> QuantumOperator:
        """零算符"""
        return QuantumOperator(np.zeros((dimension, dimension)))


# ============================================================================
# Level 3: 标准量子算符
# ============================================================================

class PauliMatrices:
    """
    Pauli矩阵: σ₀, σ₁, σ₂, σ₃
    
    量子比特的基础算符
    """
    
    @staticmethod
    def sigma_0() -> QuantumOperator:
        """单位矩阵"""
        return QuantumOperator(np.array([
            [1, 0],
            [0, 1]
        ]))
    
    @staticmethod
    def sigma_x() -> QuantumOperator:
        """Pauli X (bit flip)"""
        return QuantumOperator(np.array([
            [0, 1],
            [1, 0]
        ]))
    
    @staticmethod
    def sigma_y() -> QuantumOperator:
        """Pauli Y"""
        return QuantumOperator(np.array([
            [0, -1j],
            [1j, 0]
        ]))
    
    @staticmethod
    def sigma_z() -> QuantumOperator:
        """Pauli Z (phase flip)"""
        return QuantumOperator(np.array([
            [1, 0],
            [0, -1]
        ]))
    
    @staticmethod
    def verify_algebra() -> Dict[str, bool]:
        """验证Pauli代数"""
        sx = PauliMatrices.sigma_x()
        sy = PauliMatrices.sigma_y()
        sz = PauliMatrices.sigma_z()
        I = PauliMatrices.sigma_0()
        
        return {
            'sigma_x_squared': np.allclose(
                (sx @ sx).matrix, I.matrix
            ),
            'sigma_y_squared': np.allclose(
                (sy @ sy).matrix, I.matrix
            ),
            'sigma_z_squared': np.allclose(
                (sz @ sz).matrix, I.matrix
            ),
            'xy_commutator': np.allclose(
                sx.commutator(sy).matrix, 2j * sz.matrix
            ),
            'yz_commutator': np.allclose(
                sy.commutator(sz).matrix, 2j * sx.matrix
            ),
            'zx_commutator': np.allclose(
                sz.commutator(sx).matrix, 2j * sy.matrix
            )
        }


class QuantumGates:
    """
    量子门: Unitary算符
    
    量子计算的基本操作
    """
    
    @staticmethod
    def hadamard() -> QuantumOperator:
        """
        Hadamard门: H = (1/√2) [[1, 1], [1, -1]]
        
        创建叠加态: H|0⟩ = (|0⟩ + |1⟩)/√2
        """
        return QuantumOperator(np.array([
            [1, 1],
            [1, -1]
        ]) / np.sqrt(2))
    
    @staticmethod
    def phase(phi: float) -> QuantumOperator:
        """
        相位门: P(φ) = [[1, 0], [0, e^(iφ)]]
        """
        return QuantumOperator(np.array([
            [1, 0],
            [0, np.exp(1j * phi)]
        ]))
    
    @staticmethod
    def cnot() -> QuantumOperator:
        """
        CNOT门 (Controlled-NOT):
        |00⟩ → |00⟩
        |01⟩ → |01⟩
        |10⟩ → |11⟩
        |11⟩ → |10⟩
        """
        return QuantumOperator(np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0],
            [0, 0, 0, 1],
            [0, 0, 1, 0]
        ]))
    
    @staticmethod
    def swap() -> QuantumOperator:
        """
        SWAP门: 交换两个量子比特
        |ψφ⟩ → |φψ⟩
        """
        return QuantumOperator(np.array([
            [1, 0, 0, 0],
            [0, 0, 1, 0],
            [0, 1, 0, 0],
            [0, 0, 0, 1]
        ]))
    
    @staticmethod
    def toffoli() -> QuantumOperator:
        """
        Toffoli门 (CCNOT): 双控制非门
        """
        T = np.eye(8, dtype=complex)
        T[6:8, 6:8] = np.array([[0, 1], [1, 0]])
        return QuantumOperator(T)


class PositionMomentumOperators:
    """
    位置-动量算符
    
    连续系统的基础算符
    """
    
    @staticmethod
    def position_matrix(N: int, dx: float) -> QuantumOperator:
        """
        位置算符（离散近似）
        x̂ = x
        """
        x_values = np.arange(N) * dx - (N * dx) / 2
        return QuantumOperator(np.diag(x_values))
    
    @staticmethod
    def momentum_matrix(N: int, dx: float) -> QuantumOperator:
        """
        动量算符（离散近似）
        p̂ = -iℏ ∂/∂x
        """
        hbar = QuantumConstants.HBAR
        p_matrix = np.zeros((N, N), dtype=complex)
        
        # 有限差分近似
        for i in range(N):
            if i > 0:
                p_matrix[i, i-1] = -1j * hbar / (2 * dx)
            if i < N - 1:
                p_matrix[i, i+1] = 1j * hbar / (2 * dx)
        
        return QuantumOperator(p_matrix)
    
    @staticmethod
    def verify_heisenberg_uncertainty(
        x_op: QuantumOperator,
        p_op: QuantumOperator,
        state: QuantumState
    ) -> Dict[str, float]:
        """
        验证Heisenberg不确定性原理:
        Δx · Δp ≥ ℏ/2
        """
        delta_x = state.uncertainty(x_op)
        delta_p = state.uncertainty(p_op)
        hbar = QuantumConstants.HBAR
        
        return {
            'delta_x': delta_x,
            'delta_p': delta_p,
            'product': delta_x * delta_p,
            'lower_bound': hbar / 2,
            'satisfies_uncertainty': delta_x * delta_p >= hbar / 2 - 1e-10
        }


# ============================================================================
# Level 4: 测量理论 (Measurement Theory)
# ============================================================================

class Measurement:
    """
    量子测量
    
    基于 ±0 的理解:
    - 测量前：叠加态 (正负共存)
    - 测量中：坍缩过程 (向均衡)
    - 测量后：本征态 (确定性)
    """
    
    def __init__(self, observable: QuantumOperator):
        """
        Args:
            observable: 可观测量（Hermitian算符）
        """
        if not observable.is_hermitian():
            raise ValueError("Observable must be Hermitian")
        
        self.observable = observable
        self.eigenvalues, self.eigenvectors = observable.eigendecomposition()
    
    def measure(self, state: QuantumState, random_seed: Optional[int] = None) -> Tuple[float, QuantumState]:
        """
        执行测量（Von Neumann测量）
        
        Returns:
            measured_value: 测量结果（本征值）
            collapsed_state: 坍缩后的态（本征态）
        """
        if random_seed is not None:
            np.random.seed(random_seed)
        
        # 计算每个本征态的概率
        probabilities = []
        for i in range(len(self.eigenvalues)):
            eigenvector = self.eigenvectors[:, i]
            eigenstate = QuantumState(eigenvector, is_normalized=True)
            prob = state.measurement_probability(eigenstate)
            probabilities.append(prob)
        
        # 归一化概率
        probabilities = np.array(probabilities)
        probabilities /= probabilities.sum()
        
        # 随机选择测量结果
        outcome_index = np.random.choice(
            len(self.eigenvalues),
            p=probabilities
        )
        
        measured_value = float(np.real(self.eigenvalues[outcome_index]))
        collapsed_state = QuantumState(
            self.eigenvectors[:, outcome_index],
            is_normalized=True
        )
        
        return measured_value, collapsed_state
    
    def measurement_statistics(
        self,
        state: QuantumState,
        num_measurements: int = 1000
    ) -> Dict[str, Any]:
        """
        测量统计
        
        多次测量得到统计分布
        """
        outcomes = []
        for _ in range(num_measurements):
            value, _ = self.measure(state)
            outcomes.append(value)
        
        outcomes = np.array(outcomes)
        
        return {
            'mean': np.mean(outcomes),
            'std': np.std(outcomes),
            'theoretical_mean': float(np.real(state.expectation_value(self.observable))),
            'outcomes': outcomes
        }
    
    @staticmethod
    def projective_measurement(
        state: QuantumState,
        projectors: List[QuantumOperator]
    ) -> Tuple[int, QuantumState]:
        """
        投影测量
        
        Args:
            state: 待测量态
            projectors: 投影算符列表 {Πᵢ}
        
        Returns:
            outcome_index: 测量结果索引
            collapsed_state: 投影后的态
        """
        probabilities = []
        projected_states = []
        
        for projector in projectors:
            # Πᵢ|ψ⟩
            projected = projector @ state
            # P(i) = ⟨ψ|Πᵢ|ψ⟩
            prob = state.measurement_probability(
                QuantumState(projected.state_vector, is_normalized=True)
            )
            probabilities.append(prob)
            projected_states.append(projected)
        
        # 归一化
        probabilities = np.array(probabilities)
        probabilities /= probabilities.sum()
        
        # 随机选择
        outcome_index = np.random.choice(len(projectors), p=probabilities)
        collapsed_state = projected_states[outcome_index]
        collapsed_state.normalize()
        
        return outcome_index, collapsed_state


class POVM:
    """
    POVM (Positive Operator-Valued Measure)
    
    广义测量：允许非正交测量
    """
    
    def __init__(self, elements: List[QuantumOperator]):
        """
        Args:
            elements: POVM元素 {Eᵢ}
            满足: Σᵢ Eᵢ = I, Eᵢ ≥ 0
        """
        self.elements = elements
        
        if not self.verify_completeness():
            raise ValueError("POVM elements must sum to identity")
    
    def verify_completeness(self, tol: float = 1e-10) -> bool:
        """验证完备性: Σᵢ Eᵢ = I"""
        total = sum(elem.matrix for elem in self.elements)
        identity = np.eye(self.elements[0].dimension)
        return np.allclose(total, identity, atol=tol)
    
    def measurement_probabilities(self, state: QuantumState) -> np.ndarray:
        """
        计算测量概率: P(i) = ⟨ψ|Eᵢ|ψ⟩
        """
        probs = []
        for element in self.elements:
            prob = state.expectation_value(element)
            probs.append(np.real(prob))
        return np.array(probs)


# ============================================================================
# Level 5: 量子演化 (Quantum Evolution)
# ============================================================================

class QuantumEvolution:
    """
    量子演化
    
    - Schrödinger方程: iℏ ∂|ψ⟩/∂t = Ĥ|ψ⟩
    - 幺正演化: |ψ(t)⟩ = Û(t)|ψ(0)⟩
    """
    
    def __init__(self, hamiltonian: QuantumOperator):
        """
        Args:
            hamiltonian: 哈密顿量 Ĥ
        """
        if not hamiltonian.is_hermitian():
            raise ValueError("Hamiltonian must be Hermitian")
        
        self.hamiltonian = hamiltonian
    
    def time_evolution_operator(self, t: float, hbar: float = 1.0) -> QuantumOperator:
        """
        时间演化算符: Û(t) = exp(-iĤt/ℏ)
        """
        exponent = -1j * self.hamiltonian.matrix * t / hbar
        U_matrix = la.expm(exponent)
        return QuantumOperator(U_matrix)
    
    def evolve(
        self,
        initial_state: QuantumState,
        time: float,
        hbar: float = 1.0
    ) -> QuantumState:
        """
        演化态: |ψ(t)⟩ = Û(t)|ψ(0)⟩
        """
        U = self.time_evolution_operator(time, hbar)
        return U @ initial_state
    
    def evolve_trajectory(
        self,
        initial_state: QuantumState,
        time_points: np.ndarray,
        hbar: float = 1.0
    ) -> List[QuantumState]:
        """
        计算演化轨迹
        """
        trajectory = []
        for t in time_points:
            state_t = self.evolve(initial_state, t, hbar)
            trajectory.append(state_t)
        return trajectory
    
    def expectation_trajectory(
        self,
        initial_state: QuantumState,
        observable: QuantumOperator,
        time_points: np.ndarray,
        hbar: float = 1.0
    ) -> np.ndarray:
        """
        期望值随时间演化: ⟨Â⟩(t) = ⟨ψ(t)|Â|ψ(t)⟩
        """
        expectations = []
        for t in time_points:
            state_t = self.evolve(initial_state, t, hbar)
            exp_val = state_t.expectation_value(observable)
            expectations.append(np.real(exp_val))
        return np.array(expectations)


class SchrodingerEquation:
    """
    Schrödinger方程求解器
    
    iℏ ∂ψ/∂t = Ĥψ
    """
    
    @staticmethod
    def solve_time_dependent(
        hamiltonian_func: Callable[[float], QuantumOperator],
        initial_state: QuantumState,
        time_points: np.ndarray,
        hbar: float = 1.0
    ) -> List[QuantumState]:
        """
        求解含时Schrödinger方程
        
        Args:
            hamiltonian_func: 时间依赖的哈密顿量 H(t)
            initial_state: 初态
            time_points: 时间点数组
            hbar: 约化普朗克常数
        """
        def deriv(psi_flat, t):
            # 重构复数态向量
            dim = len(psi_flat) // 2
            psi = psi_flat[:dim] + 1j * psi_flat[dim:]
            
            # 计算 -iĤψ/ℏ
            H_t = hamiltonian_func(t)
            dpsi = -1j * H_t.matrix @ psi / hbar
            
            # 展平为实数数组
            dpsi_flat = np.concatenate([np.real(dpsi), np.imag(dpsi)])
            return dpsi_flat
        
        # 初态展平
        psi0 = initial_state.state_vector
        psi0_flat = np.concatenate([np.real(psi0), np.imag(psi0)])
        
        # 求解ODE
        solution = odeint(deriv, psi0_flat, time_points)
        
        # 重构量子态
        states = []
        dim = len(psi0)
        for psi_flat in solution:
            psi = psi_flat[:dim] + 1j * psi_flat[dim:]
            states.append(QuantumState(psi, is_normalized=True))
        
        return states


# ============================================================================
# Level 6: 量子纠缠 (Quantum Entanglement)
# ============================================================================

class EntangledState:
    """
    纠缠态
    
    不可分解为直积态: |ψ⟩ ≠ |φ₁⟩ ⊗ |φ₂⟩
    """
    
    def __init__(self, state_vector: np.ndarray, subsystem_dims: Tuple[int, ...]):
        """
        Args:
            state_vector: 组合系统的态向量
            subsystem_dims: 各子系统的维数
        """
        self.state = QuantumState(state_vector)
        self.subsystem_dims = subsystem_dims
        self.total_dim = np.prod(subsystem_dims)
        
        if len(state_vector) != self.total_dim:
            raise ValueError("State dimension mismatch")
    
    @staticmethod
    def bell_state(which: str = 'phi_plus') -> EntangledState:
        """
        Bell态（最大纠缠态）
        
        |Φ⁺⟩ = (|00⟩ + |11⟩)/√2
        |Φ⁻⟩ = (|00⟩ - |11⟩)/√2
        |Ψ⁺⟩ = (|01⟩ + |10⟩)/√2
        |Ψ⁻⟩ = (|01⟩ - |10⟩)/√2
        """
        bell_states = {
            'phi_plus': np.array([1, 0, 0, 1]) / np.sqrt(2),
            'phi_minus': np.array([1, 0, 0, -1]) / np.sqrt(2),
            'psi_plus': np.array([0, 1, 1, 0]) / np.sqrt(2),
            'psi_minus': np.array([0, 1, -1, 0]) / np.sqrt(2)
        }
        
        if which not in bell_states:
            raise ValueError(f"Unknown Bell state: {which}")
        
        return EntangledState(bell_states[which], (2, 2))
    
    def partial_trace(self, subsystem: int) -> np.ndarray:
        """
        偏迹: ρ_A = Tr_B(|ψ⟩⟨ψ|)
        
        Args:
            subsystem: 要保留的子系统索引（0或1）
        """
        # 构造密度矩阵
        psi = self.state.state_vector
        rho = np.outer(psi, np.conj(psi))
        
        # 重塑为张量形式
        dims = self.subsystem_dims
        rho_tensor = rho.reshape(dims + dims)
        
        # 执行偏迹
        if subsystem == 0:
            # 保留第一个子系统
            axes = (1, 3)
        else:
            # 保留第二个子系统
            axes = (0, 2)
        
        rho_reduced = np.trace(rho_tensor, axis1=axes[0], axis2=axes[1])
        
        return rho_reduced
    
    def entropy_of_entanglement(self) -> float:
        """
        纠缠熵: S = -Tr(ρ_A log ρ_A)
        
        度量纠缠程度
        """
        rho_A = self.partial_trace(0)
        
        # 计算本征值
        eigenvalues = np.linalg.eigvalsh(rho_A)
        eigenvalues = eigenvalues[eigenvalues > 1e-10]  # 过滤接近零的值
        
        # S = -Σ λᵢ log₂ λᵢ
        entropy = -np.sum(eigenvalues * np.log2(eigenvalues))
        
        return float(entropy)
    
    def concurrence(self) -> float:
        """
        Concurrence: 两量子比特纠缠度量
        
        C = 0: 可分离态
        C = 1: 最大纠缠
        """
        if self.subsystem_dims != (2, 2):
            raise ValueError("Concurrence only defined for two qubits")
        
        psi = self.state.state_vector
        
        # 构造密度矩阵
        rho = np.outer(psi, np.conj(psi))
        
        # Pauli Y ⊗ Y
        sigma_y = np.array([[0, -1j], [1j, 0]])
        Y = np.kron(sigma_y, sigma_y)
        
        # ρ̃ = (σy ⊗ σy) ρ* (σy ⊗ σy)
        rho_tilde = Y @ np.conj(rho) @ Y
        
        # R = ρ ρ̃
        R = rho @ rho_tilde
        
        # 计算R的本征值
        eigenvalues = np.linalg.eigvalsh(R)
        eigenvalues = np.sqrt(np.maximum(eigenvalues, 0))
        eigenvalues = np.sort(eigenvalues)[::-1]
        
        # C = max(0, λ₁ - λ₂ - λ₃ - λ₄)
        concurrence = max(0, eigenvalues[0] - eigenvalues[1] - eigenvalues[2] - eigenvalues[3])
        
        return float(concurrence)
    
    def is_separable(self, tol: float = 1e-6) -> bool:
        """判断是否为可分离态"""
        entropy = self.entropy_of_entanglement()
        return entropy < tol


class QuantumTeleportation:
    """
    量子隐形传态
    
    基于纠缠和经典通信传输量子态
    """
    
    @staticmethod
    def teleport(
        state_to_teleport: QuantumState,
        shared_entangled_state: EntangledState
    ) -> Tuple[List[int], QuantumState]:
        """
        执行量子隐形传态
        
        Args:
            state_to_teleport: 待传输的态 |ψ⟩
            shared_entangled_state: Alice和Bob共享的纠缠态
        
        Returns:
            classical_bits: Alice的测量结果（经典信息）
            teleported_state: Bob接收到的态
        """
        # 简化实现：使用Bell态
        # |ψ⟩|Φ⁺⟩ → Bell测量 → 经典通信 → Bob的幺正操作
        
        # Alice的Bell测量（简化）
        measurement_results = [np.random.randint(0, 2), np.random.randint(0, 2)]
        
        # Bob根据测量结果应用幺正操作
        if measurement_results == [0, 0]:
            U = QuantumOperator(np.eye(2))
        elif measurement_results == [0, 1]:
            U = PauliMatrices.sigma_x()
        elif measurement_results == [1, 0]:
            U = PauliMatrices.sigma_z()
        else:
            U = PauliMatrices.sigma_x() @ PauliMatrices.sigma_z()
        
        teleported_state = U @ state_to_teleport
        
        return measurement_results, teleported_state


# ============================================================================
# Level 7: 量子系统实例
# ============================================================================

class TwoLevelSystem:
    """
    二能级系统（量子比特）
    
    最简单的量子系统
    """
    
    @staticmethod
    def ground_state() -> QuantumState:
        """基态 |0⟩"""
        return QuantumState([1, 0])
    
    @staticmethod
    def excited_state() -> QuantumState:
        """激发态 |1⟩"""
        return QuantumState([0, 1])
    
    @staticmethod
    def superposition(theta: float, phi: float) -> QuantumState:
        """
        一般叠加态: |ψ⟩ = cos(θ/2)|0⟩ + e^(iφ)sin(θ/2)|1⟩
        
        Bloch球表示
        """
        return QuantumState([
            np.cos(theta / 2),
            np.exp(1j * phi) * np.sin(theta / 2)
        ])
    
    @staticmethod
    def rabi_oscillation(
        omega: float,
        time_points: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Rabi振荡
        
        在共振外场驱动下，二能级系统的周期振荡
        """
        # Hamiltonian: H = ℏω/2 σ_x
        H = 0.5 * omega * PauliMatrices.sigma_x()
        evolution = QuantumEvolution(H)
        
        initial_state = TwoLevelSystem.ground_state()
        
        # 计算激发态布居数
        excited_populations = []
        ground_populations = []
        
        for t in time_points:
            state_t = evolution.evolve(initial_state, t, hbar=1.0)
            
            # |⟨1|ψ(t)⟩|²
            excited_prob = state_t.measurement_probability(
                TwoLevelSystem.excited_state()
            )
            excited_populations.append(excited_prob)
            ground_populations.append(1 - excited_prob)
        
        return np.array(ground_populations), np.array(excited_populations)


class HarmonicOscillator:
    """
    量子谐振子
    
    H = (p²/2m) + (mω²x²/2)
    """
    
    def __init__(self, mass: float, omega: float, N: int = 50):
        """
        Args:
            mass: 质量
            omega: 角频率
            N: 离散化维数
        """
        self.mass = mass
        self.omega = omega
        self.N = N
        
        # 构造哈密顿量
        self.hamiltonian = self._build_hamiltonian()
    
    def _build_hamiltonian(self) -> QuantumOperator:
        """构造离散哈密顿量"""
        # 使用能量本征态基
        # H = ℏω(â†â + 1/2)
        hbar = 1.0  # 自然单位
        
        # 能量本征值
        eigenvalues = hbar * self.omega * (np.arange(self.N) + 0.5)
        H_matrix = np.diag(eigenvalues)
        
        return QuantumOperator(H_matrix)
    
    def energy_eigenstate(self, n: int) -> QuantumState:
        """
        能量本征态 |n⟩
        
        Args:
            n: 量子数
        """
        if n >= self.N:
            raise ValueError(f"n must be < {self.N}")
        
        state_vector = np.zeros(self.N, dtype=complex)
        state_vector[n] = 1.0
        
        return QuantumState(state_vector, is_normalized=True)
    
    def coherent_state(self, alpha: complex) -> QuantumState:
        """
        相干态 |α⟩
        
        最接近经典态的量子态
        |α⟩ = exp(-|α|²/2) Σ (αⁿ/√n!) |n⟩
        """
        state_vector = np.zeros(self.N, dtype=complex)
        
        normalization = np.exp(-np.abs(alpha)**2 / 2)
        
        for n in range(self.N):
            coefficient = normalization * (alpha**n) / np.sqrt(np.math.factorial(n))
            state_vector[n] = coefficient
        
        return QuantumState(state_vector, is_normalized=True)


# ============================================================================
# 量子力学完整演示
# ============================================================================

class QuantumMechanicsDemo:
    """量子力学演示系统"""
    
    @staticmethod
    def demonstrate_superposition():
        """演示叠加原理"""
        print("\n" + "=" * 70)
        print("叠加原理演示")
        print("=" * 70)
        
        # 创建叠加态
        psi = TwoLevelSystem.superposition(np.pi/4, 0)
        print(f"叠加态: |ψ⟩ = {psi.state_vector}")
        
        # 测量概率
        ground = TwoLevelSystem.ground_state()
        excited = TwoLevelSystem.excited_state()
        
        p_ground = psi.measurement_probability(ground)
        p_excited = psi.measurement_probability(excited)
        
        print(f"测量得到 |0⟩ 的概率: {p_ground:.4f}")
        print(f"测量得到 |1⟩ 的概率: {p_excited:.4f}")
        print(f"概率和: {p_ground + p_excited:.4f}")
    
    @staticmethod
    def demonstrate_uncertainty():
        """演示不确定性原理"""
        print("\n" + "=" * 70)
        print("不确定性原理演示")
        print("=" * 70)
        
        # 创建态
        psi = TwoLevelSystem.superposition(np.pi/3, 0)
        
        # Pauli算符
        sx = PauliMatrices.sigma_x()
        sy = PauliMatrices.sigma_y()
        sz = PauliMatrices.sigma_z()
        
        # 计算不确定度
        delta_x = psi.uncertainty(sx)
        delta_y = psi.uncertainty(sy)
        delta_z = psi.uncertainty(sz)
        
        print(f"Δσx = {delta_x:.4f}")
        print(f"Δσy = {delta_y:.4f}")
        print(f"Δσz = {delta_z:.4f}")
        
        # 验证不确定性关系
        print(f"\n不确定性乘积:")
        print(f"Δσx · Δσy = {delta_x * delta_y:.4f} ≥ |⟨[σx,σy]⟩|/2")
    
    @staticmethod
    def demonstrate_evolution():
        """演示量子演化"""
        print("\n" + "=" * 70)
        print("量子演化演示 (Rabi振荡)")
        print("=" * 70)
        
        omega = 1.0
        time_points = np.linspace(0, 4*np.pi/omega, 100)
        
        ground_pop, excited_pop = TwoLevelSystem.rabi_oscillation(
            omega, time_points
        )
        
        print(f"初态: |0⟩")
        print(f"哈密顿量: H = (ℏω/2)σx")
        print(f"\n在 t=π/(2ω) 时:")
        idx = len(time_points) // 4
        print(f"  P(|0⟩) = {ground_pop[idx]:.4f}")
        print(f"  P(|1⟩) = {excited_pop[idx]:.4f}")
    
    @staticmethod
    def demonstrate_entanglement():
        """演示量子纠缠"""
        print("\n" + "=" * 70)
        print("量子纠缠演示")
        print("=" * 70)
        
        # Bell态
        bell = EntangledState.bell_state('phi_plus')
        print(f"Bell态 |Φ⁺⟩ = (|00⟩ + |11⟩)/√2")
        print(f"态向量: {bell.state.state_vector}")
        
        # 纠缠熵
        entropy = bell.entropy_of_entanglement()
        print(f"\n纠缠熵: S = {entropy:.4f}")
        print(f"(S=1 表示最大纠缠)")
        
        # Concurrence
        conc = bell.concurrence()
        print(f"Concurrence: C = {conc:.4f}")
        print(f"(C=1 表示最大纠缠)")
        
        # 可分离性
        is_sep = bell.is_separable()
        print(f"\n是否可分离: {is_sep}")
    
    @staticmethod
    def demonstrate_measurement():
        """演示量子测量"""
        print("\n" + "=" * 70)
        print("量子测量演示")
        print("=" * 70)
        
        # 创建叠加态
        psi = TwoLevelSystem.superposition(np.pi/4, 0)
        print(f"初态: |ψ⟩")
        
        # 在 σz 基下测量
        measurement = Measurement(PauliMatrices.sigma_z())
        
        # 多次测量统计
        stats = measurement.measurement_statistics(psi, num_measurements=1000)
        
        print(f"\n测量统计 (1000次):")
        print(f"平均值: {stats['mean']:.4f}")
        print(f"标准差: {stats['std']:.4f}")
        print(f"理论期望: {stats['theoretical_mean']:.4f}")


def main():
    """主函数：完整演示"""
    
    print("=" * 70)
    print("DaoMath Quantum Mechanics")
    print("从 ±0 构造量子理论")
    print("=" * 70)
    
    # 1. 基础验证
    print("\n" + "=" * 70)
    print("Pauli矩阵验证")
    print("=" * 70)
    pauli_check = PauliMatrices.verify_algebra()
    for prop, satisfied in pauli_check.items():
        print(f"  • {prop}: {satisfied}")
    
    # 2. 叠加原理
    QuantumMechanicsDemo.demonstrate_superposition()
    
    # 3. 不确定性原理
    QuantumMechanicsDemo.demonstrate_uncertainty()
    
    # 4. 量子演化
    QuantumMechanicsDemo.demonstrate_evolution()
    
    # 5. 量子纠缠
    QuantumMechanicsDemo.demonstrate_entanglement()
    
    # 6. 量子测量
    QuantumMechanicsDemo.demonstrate_measurement()
    
    print("\n" + "=" * 70)
    print("量子力学理论构造完成！")
    print("=" * 70)


if __name__ == "__main__":
    main()
